
import os
import cv2 as cv
import shutil
from xml.dom.minidom import Document

from PIL import Image
from yolo import YOLO


def detect_img(yolo):
    while True:
        img = input('Input image filename:')
        try:
            image = Image.open(img)
        except:
            print('Open Error! Try again!')
            continue
        else:
            r_image, label, left, top, right, bottom = yolo.detect_image(image)
            r_image.show()
    yolo.close_session()


def main():
    img_dir = './data/images'        # 图片目录，读取tfrecord文件解压的图片
    res_dir = './data/detec_image'  # 输出目录，保存绘制bbox的图片
    err_dir = './data/err_image'    # 错误目录，保存未识别到特征的图片
    xml_dir = './data/annotations/xmls'          # xml目录，保存xml文件

    image_names = os.listdir(img_dir)  # 图片名称列表
    num = 0
    yolo = YOLO()
    for name in image_names:
        if name is None:
            print("No found image!")
            return 0
        num += 1
        # print("\nThe {}th Image read OK!!!：{}".format(num, name))

        # 读取图片名称
        if name[-15] == '_':
            img_label = int(name[-14:-12])
        else:
            img_label = int(name[-15:-12])

        image_file = img_dir + '/' + name
        image = Image.open(image_file)
        # image.show()

        # 返回图片分类和bbox信息
        try:
            r_image, label, score, left, top, right, bottom, flag = yolo.detect_image(image)
        except:
            print("Not found feature in this picture")
            image.save(err_dir + '/' + str(num) + '_' + str(img_label) + '.jpg')
            continue
        else:
            # imgfile = res_dir + '/' + str(num) + '_' + name
            imgfile = res_dir + '/' + name   # 绘制了bbox的图片名称
            # r_image.save(imgfile)
            # print("写入成功!")

            # if flag == True:
            #     # 返回多个
            #     print("图像中存在多个特征，筛选后仅保留了一个最合格的特征！")

            # xml文件名
            # xmlfile = xml_dir + '/' + str(num) + '_' + name.split('.')[0] + '.xml'
            xmlfile = xml_dir + '/' + name.split('.')[0] + '.xml'
            # print("----->xmlfile:",xmlfile)

            # 创建doc
            doc = Document()
            # 创建根节点
            root = doc.createElement("annotation")
            doc.appendChild(root)
            nodes = {0: root}

            node = doc.createElement("filename")
            node.appendChild(doc.createTextNode(name))
            nodes[0].appendChild(node)
            nodes[1] = node

            node_path = doc.createElement("path")
            node_path.appendChild(doc.createTextNode(image_file))
            nodes[0].appendChild(node_path)
            nodes[1] = node_path

            # 添加size节点
            node1 = doc.createElement("size")
            nodes[0].appendChild(node1)
            nodes[1] = node1
            param = name.split(".")[0].split("_")
            
            # 添加宽度节点
            node_width_size = doc.createElement("width")
            node_width_size.appendChild(doc.createTextNode(param[3]))
            nodes[1].appendChild(node_width_size)
            nodes[2] = node_width_size
            # 添加高度节点
            node_height_size = doc.createElement("height")
            node_height_size.appendChild(doc.createTextNode(param[4]))
            nodes[1].appendChild(node_height_size)
            nodes[2] = node_height_size

            node_depth = doc.createElement("depth")
            node_depth.appendChild(doc.createTextNode("3"))
            nodes[1].appendChild(node_depth)
            nodes[2] = node_depth

            # for i in range(len(label)):
            node1_object = doc.createElement("object")
            nodes[0].appendChild(node1_object)
            nodes[1] = node1_object

            node_class_object = doc.createElement("name")
            node_class_object.appendChild(doc.createTextNode(str(name.split(".")[0].split("_")[2])))

            nodes[1].appendChild(node_class_object)
            nodes[2] = node_class_object

            node_pose = doc.createElement("pose")
            node_pose.appendChild(doc.createTextNode("Unspecified"))
            nodes[1].appendChild(node_pose)
            nodes[2] = node_pose

            node_truncated = doc.createElement("truncated")
            node_truncated.appendChild(doc.createTextNode("0"))
            nodes[1].appendChild(node_truncated)
            nodes[2] = node_truncated

            node_difficult = doc.createElement("difficult")
            node_difficult.appendChild(doc.createTextNode("0"))
            nodes[1].appendChild(node_difficult)
            nodes[2] = node_difficult

            node_bndbox = doc.createElement("bndbox")

            nodes[1].appendChild(node_bndbox)
            nodes[2] = node_bndbox
            # bbbox = bboxes[i]

            node_xmin = doc.createElement("xmin")
            node_xmin.appendChild(doc.createTextNode(str(left)))
            nodes[2].appendChild(node_xmin)
            nodes[3] = node_xmin

            node_ymin = doc.createElement("ymin")
            node_ymin.appendChild(doc.createTextNode(str(top)))
            nodes[2].appendChild(node_ymin)
            nodes[3] = node_ymin

            node_xmax = doc.createElement("xmax")
            node_xmax.appendChild(doc.createTextNode(str(right)))
            nodes[2].appendChild(node_xmax)
            nodes[3] = node_xmax

            node_ymax = doc.createElement("ymax")
            node_ymax.appendChild(doc.createTextNode(str(bottom)))
            nodes[2].appendChild(node_ymax)
            nodes[3] = node_ymax

            with open(xmlfile, 'w') as f:
                f.write(doc.toprettyxml(indent='\t'))

    yolo.close_session()


# FLAGS = None

if __name__ == '__main__':
    main()

    # # class YOLO defines the default value, so suppress any default here
    # parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS)
    # '''
    # Command line options
    # '''
    # parser.add_argument(
    #     '--model', type=str,
    #     help='path to model weight file, default ' + YOLO.get_defaults("model_path")
    # )
    #
    # parser.add_argument(
    #     '--anchors', type=str,
    #     help='path to anchor definitions, default ' + YOLO.get_defaults("anchors_path")
    # )
    #
    # parser.add_argument(
    #     '--classes', type=str,
    #     help='path to class definitions, default ' + YOLO.get_defaults("classes_path")
    # )
    #
    # parser.add_argument(
    #     '--gpu_num', type=int,
    #     help='Number of GPU to use, default ' + str(YOLO.get_defaults("gpu_num"))
    # )
    #
    # parser.add_argument(
    #     '--image', default=False, action="store_true",
    #     help='Image detection mode, will ignore all positional arguments'
    # )
    # '''
    # Command line positional arguments -- for video detection mode
    # '''
    # parser.add_argument(
    #     "--input", nargs='?', type=str, required=False, default='./path2your_video',
    #     help="Video input path"
    # )
    #
    # parser.add_argument(
    #     "--output", nargs='?', type=str, default="",
    #     help="[Optional] Video output path"
    # )
    #
    # FLAGS = parser.parse_args()
    #
    # if FLAGS.image:
    #     """
    #     Image detection mode, disregard any remaining command line arguments
    #     """
    #     print("Image detection mode")
    #     if "input" in FLAGS:
    #         print(" Ignoring remaining command line arguments: " + FLAGS.input + "," + FLAGS.output)
    #     detect_img(YOLO(**vars(FLAGS)))
    # elif "input" in FLAGS:
    #     detect_video(YOLO(**vars(FLAGS)), FLAGS.input, FLAGS.output)
    # else:
    #     print("Must specify at least video_input_path.  See usage with --help.")
